import numpy as np
import os
import array
import cv2
import shutil
from tqdm import tqdm
import time


def generate_and_save_image(input_dir, output_dir, filename, str=""):
    """
        生成图片
    """
    if str != "":
        temp = str.split("_")[1]
        output_dir += temp
        output_dir += "/"
    out_file = os.path.splitext(str+os.path.basename(filename))[0] + '.png'
    out_file_full = output_dir + out_file
    input_file_path = os.path.join(input_dir, filename)

    f = open(input_file_path, 'rb')
    ln = os.path.getsize(input_file_path)
    width = 256
    rem = ln % width

    a = array.array("B")
    a.fromfile(f, ln - rem)
    f.close()
    g = np.reshape(a, (len(a) // width, width))
    g = np.uint8(g)
    img_resized = cv2.resize(g, (105, 105), interpolation=cv2.INTER_CUBIC)
    cv2.imwrite(out_file_full, img_resized)


def convert_bin_to_img(input_dir, output_dir):
    """
        输入为存放二进制文件的地方
        输出为图片存放的地方
    """
    a = time.time()
    try:
        shutil.rmtree(output_dir)
    except Exception as e:
        print(e)
    os.mkdir(output_dir)

    count = 0
    files = os.listdir(input_dir)
    errorName = []
    for filename in tqdm(files):
        try:
            generate_and_save_image(input_dir, output_dir+"/", filename)
            count += 1
        except Exception as e:
            print(e)
            errorName.append(filename)
    print("{}个样本被转换成图片".format(count))
    print("{}失败的样本名如下：".format(len(errorName)))
    print(errorName)
    b = time.time()
    print("生成耗时："+str(b-a))


if __name__ == "__main__":
    # convert_bin_to_img("/home/cyw/projects/data/functionsim/allorigindata","/home/cyw/projects/function_sim_project/all_data/sampleDatas/img")
    convert_bin_to_img("/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/dataTrans",
                       "/mnt/mydisk1/chenyongwei/malware/BODMAS_dataset/img")
    # success 2614   error 7
